"""
数据库结构智能增强器
使用AI模型对数据库结构进行智能分析和补全
"""

import logging
from typing import Dict, List, Any, Optional, Tuple
from concurrent.futures import ThreadPoolExecutor, as_completed
import threading
import re
import json

from .model_interface import ai_model_manager
from src.database.schema_extractor import DatabaseSchema, TableInfo, ColumnInfo
from config.ai_config import PromptTemplates


class SchemaEnhancer:
    """数据库结构智能增强器"""
    
    def __init__(self):
        self.logger = logging.getLogger(__name__)
        self.lock = threading.RLock()
        self.enhancement_history: Dict[str, Dict] = {}
    
    def enhance_database_schema(self, schema: DatabaseSchema, 
                              enhance_table_comments: bool = True,
                              enhance_column_comments: bool = True,
                              analyze_relationships: bool = True,
                              max_workers: int = 3) -> DatabaseSchema:
        """
        智能增强数据库结构
        
        Args:
            schema: 原始数据库结构
            enhance_table_comments: 是否增强表注释
            enhance_column_comments: 是否增强字段注释
            analyze_relationships: 是否分析表关系
            max_workers: 最大并发工作线程数
        """
        try:
            self.logger.info(f"开始智能增强数据库结构: {schema.database_name}")
            
            # 创建增强历史记录
            enhancement_key = f"{schema.config_name}_{schema.database_name}"
            self.enhancement_history[enhancement_key] = {
                "enhanced_tables": [],
                "enhanced_columns": [],
                "analyzed_relationships": [],
                "enhancement_time": None
            }
            
            # 并发处理各个schema
            with ThreadPoolExecutor(max_workers=max_workers) as executor:
                future_to_schema = {}
                
                for schema_name, tables in schema.schemas.items():
                    future = executor.submit(
                        self._enhance_schema_tables,
                        tables, schema_name, schema.database_name,
                        enhance_table_comments, enhance_column_comments, analyze_relationships
                    )
                    future_to_schema[future] = schema_name
                
                # 收集结果
                for future in as_completed(future_to_schema):
                    schema_name = future_to_schema[future]
                    try:
                        enhanced_tables = future.result()
                        if enhanced_tables:
                            schema.schemas[schema_name] = enhanced_tables
                        
                        self.logger.info(f"Schema {schema_name} 增强完成")
                        
                    except Exception as e:
                        self.logger.error(f"增强schema {schema_name} 失败: {str(e)}")
            
            # 设置增强时间
            from datetime import datetime
            self.enhancement_history[enhancement_key]["enhancement_time"] = datetime.now().isoformat()
            
            self.logger.info(f"数据库结构增强完成: {schema.database_name}")
            return schema
            
        except Exception as e:
            self.logger.error(f"数据库结构增强失败: {str(e)}")
            return schema
    
    def _enhance_schema_tables(self, tables: List[TableInfo], schema_name: str, database_name: str,
                             enhance_table_comments: bool, enhance_column_comments: bool,
                             analyze_relationships: bool) -> List[TableInfo]:
        """增强指定schema的表"""
        enhanced_tables = []
        
        for table in tables:
            try:
                enhanced_table = self._enhance_single_table(
                    table, tables, schema_name, database_name,
                    enhance_table_comments, enhance_column_comments, analyze_relationships
                )
                enhanced_tables.append(enhanced_table)
                
            except Exception as e:
                self.logger.error(f"增强表 {table.table_name} 失败: {str(e)}")
                enhanced_tables.append(table)  # 保留原始表信息
        
        return enhanced_tables
    
    def _enhance_single_table(self, table: TableInfo, all_tables: List[TableInfo],
                            schema_name: str, database_name: str,
                            enhance_table_comments: bool, enhance_column_comments: bool,
                            analyze_relationships: bool) -> TableInfo:
        """增强单个表"""
        enhanced_table = table
        
        # 1. 增强表注释
        if enhance_table_comments and not (table.table_comment and table.table_comment.strip()):
            enhanced_comment = self._generate_table_comment(table, all_tables)
            if enhanced_comment:
                enhanced_table.table_comment = enhanced_comment
        
        # 2. 增强字段注释
        if enhance_column_comments:
            enhanced_columns = []
            for column in table.columns:
                # 强制为所有空白字段生成注释
                force_generate = not (column.column_comment and column.column_comment.strip())
                enhanced_column = self._enhance_column_comment(column, table, all_tables, force_generate)
                enhanced_columns.append(enhanced_column)
            enhanced_table.columns = enhanced_columns
        
        # 3. 分析表关系（添加到表注释的末尾）
        if analyze_relationships:
            relationship_analysis = self._analyze_table_relationships(table, all_tables)
            if relationship_analysis:
                if enhanced_table.table_comment:
                    enhanced_table.table_comment += f" | 关系分析: {relationship_analysis}"
                else:
                    enhanced_table.table_comment = f"关系分析: {relationship_analysis}"
        
        return enhanced_table
    
    def _generate_table_comment(self, table: TableInfo, all_tables: List[TableInfo]) -> Optional[str]:
        """生成表注释"""
        try:
            # 准备字段信息
            fields_info = []
            for col in table.columns[:10]:  # 只取前10个字段避免prompt过长
                field_desc = f"- {col.column_name}({col.data_type})"
                if col.is_primary_key:
                    field_desc += " [主键]"
                if col.is_foreign_key:
                    field_desc += " [外键]"
                if col.column_comment:
                    field_desc += f" - {col.column_comment}"
                fields_info.append(field_desc)
            
            # 构建提示词
            prompt = PromptTemplates.TABLE_COMMENT_TEMPLATE.format(
                table_name=table.table_name,
                fields_info="\n".join(fields_info)
            )
            
            # 调用AI模型
            response = ai_model_manager.generate_text(prompt, max_tokens=100, temperature=0.3)
            
            if response:
                # 清理响应
                comment = self._clean_ai_response(response)
                if comment and len(comment) <= 200:  # 限制注释长度
                    # 添加AI生成标记
                    ai_comment = f"{comment}[AI生成]"
                    self.logger.info(f"为表 {table.table_name} 生成注释: {comment}")
                    return ai_comment
            
            return None
            
        except Exception as e:
            self.logger.error(f"生成表注释失败 {table.table_name}: {str(e)}")
            return None
    
    def _enhance_column_comment(self, column: ColumnInfo, table: TableInfo, 
                              all_tables: List[TableInfo], force_generate: bool = True) -> ColumnInfo:
        """增强字段注释"""
        if not force_generate and column.column_comment and column.column_comment.strip():
            return column  # 已有注释，且非强制模式，不需要增强
        
        try:
            # 如果已有注释且非强制生成模式，则跳过
            if not force_generate:
                return column
            
            # 准备其他字段信息
            other_fields = [f"{col.column_name}({col.data_type})" 
                          for col in table.columns[:5] if col.column_name != column.column_name]
            
            # 构建提示词
            prompt = PromptTemplates.FIELD_COMMENT_TEMPLATE.format(
                table_name=table.table_name,
                field_name=column.column_name,
                field_type=column.data_type,
                is_nullable="可空" if column.nullable else "非空",
                default_value=column.default_value or "无",
                other_fields=", ".join(other_fields)
            )
            
            # 调用AI模型
            response = ai_model_manager.generate_text(prompt, max_tokens=50, temperature=0.3)
            
            if response:
                # 清理响应
                comment = self._clean_ai_response(response)
                if comment and len(comment) <= 100:  # 限制注释长度
                    # 添加AI生成标记
                    ai_comment = f"{comment}[AI生成]"
                    enhanced_column = ColumnInfo(
                        column_name=column.column_name,
                        data_type=column.data_type,
                        nullable=column.nullable,
                        default_value=column.default_value,
                        is_primary_key=column.is_primary_key,
                        is_foreign_key=column.is_foreign_key,
                        auto_increment=column.auto_increment,
                        column_comment=ai_comment,
                        character_maximum_length=column.character_maximum_length,
                        numeric_precision=column.numeric_precision,
                        numeric_scale=column.numeric_scale
                    )
                    
                    self.logger.info(f"为字段 {table.table_name}.{column.column_name} 生成注释: {comment}")
                    return enhanced_column
            
            # 如果AI生成失败，为空白字段提供基于字段名的简单注释
            if not (column.column_comment and column.column_comment.strip()):
                fallback_comment = self._generate_fallback_comment(column)
                if fallback_comment:
                    # 添加AI生成标记
                    ai_fallback_comment = f"{fallback_comment}[AI生成]"
                    enhanced_column = ColumnInfo(
                        column_name=column.column_name,
                        data_type=column.data_type,
                        nullable=column.nullable,
                        default_value=column.default_value,
                        is_primary_key=column.is_primary_key,
                        is_foreign_key=column.is_foreign_key,
                        auto_increment=column.auto_increment,
                        column_comment=ai_fallback_comment,
                        character_maximum_length=column.character_maximum_length,
                        numeric_precision=column.numeric_precision,
                        numeric_scale=column.numeric_scale
                    )
                    
                    self.logger.info(f"为字段 {table.table_name}.{column.column_name} 生成备选注释: {fallback_comment}")
                    return enhanced_column
            
            return column
            
        except Exception as e:
            self.logger.error(f"增强字段注释失败 {table.table_name}.{column.column_name}: {str(e)}")
            # 尝试生成备选注释
            if not (column.column_comment and column.column_comment.strip()):
                fallback_comment = self._generate_fallback_comment(column)
                if fallback_comment:
                    # 添加AI生成标记
                    ai_fallback_comment = f"{fallback_comment}[AI生成]"
                    enhanced_column = ColumnInfo(
                        column_name=column.column_name,
                        data_type=column.data_type,
                        nullable=column.nullable,
                        default_value=column.default_value,
                        is_primary_key=column.is_primary_key,
                        is_foreign_key=column.is_foreign_key,
                        auto_increment=column.auto_increment,
                        column_comment=ai_fallback_comment,
                        character_maximum_length=column.character_maximum_length,
                        numeric_precision=column.numeric_precision,
                        numeric_scale=column.numeric_scale
                    )
                    return enhanced_column
            return column
    
    def _analyze_table_relationships(self, table: TableInfo, all_tables: List[TableInfo]) -> Optional[str]:
        """分析表关系"""
        try:
            # 准备相关表信息
            related_tables = []
            table_names = [t.table_name for t in all_tables]
            
            # 查找相关表（通过外键、相似命名等）
            for other_table in all_tables:
                if other_table.table_name == table.table_name:
                    continue
                
                # 检查是否有外键关系
                has_fk_relation = False
                for fk in table.foreign_keys:
                    if fk.referenced_table == other_table.table_name:
                        has_fk_relation = True
                        break
                
                for fk in other_table.foreign_keys:
                    if fk.referenced_table == table.table_name:
                        has_fk_relation = True
                        break
                
                if has_fk_relation:
                    related_tables.append(f"{other_table.table_name}({other_table.table_comment or '无注释'})")
            
            # 如果相关表太少，添加一些命名相似的表
            if len(related_tables) < 3:
                for other_table in all_tables[:10]:  # 限制数量
                    if other_table.table_name != table.table_name:
                        related_tables.append(f"{other_table.table_name}({other_table.table_comment or '无注释'})")
                    if len(related_tables) >= 5:
                        break
            
            # 准备字段信息
            fields_info = []
            for col in table.columns[:8]:  # 限制字段数量
                field_desc = f"- {col.column_name}({col.data_type})"
                if col.is_primary_key:
                    field_desc += " [主键]"
                if col.is_foreign_key:
                    field_desc += " [外键]"
                fields_info.append(field_desc)
            
            # 构建提示词
            prompt = PromptTemplates.TABLE_RELATIONSHIP_TEMPLATE.format(
                table_name=table.table_name,
                table_comment=table.table_comment or "无注释",
                fields_info="\n".join(fields_info),
                related_tables="\n".join(related_tables[:5])  # 最多5个相关表
            )
            
            # 调用AI模型
            response = ai_model_manager.generate_text(prompt, max_tokens=150, temperature=0.4)
            
            if response:
                analysis = self._clean_ai_response(response)
                if analysis and len(analysis) <= 300:
                    self.logger.info(f"为表 {table.table_name} 生成关系分析")
                    return analysis
            
            return None
            
        except Exception as e:
            self.logger.error(f"分析表关系失败 {table.table_name}: {str(e)}")
            return None
    
    def _clean_ai_response(self, response: str) -> str:
        """清理AI响应文本"""
        if not response:
            return ""
        
        # 移除常见的无用前缀和后缀
        response = response.strip()
        
        # 移除引号
        if response.startswith('"') and response.endswith('"'):
            response = response[1:-1]
        if response.startswith("'") and response.endswith("'"):
            response = response[1:-1]
        
        # 移除常见的AI回复模式
        patterns_to_remove = [
            r'^(这个表|该表|此表|表)\s*',
            r'^(这个字段|该字段|此字段|字段)\s*',
            r'^(根据分析|分析结果|综合分析)\s*[:,：]\s*',
            r'^(答案|回答|结果)\s*[:,：]\s*',
            r'\s*(。|\.)\s*$'
        ]
        
        for pattern in patterns_to_remove:
            response = re.sub(pattern, '', response, flags=re.IGNORECASE)
        
        return response.strip()
    
    def _generate_fallback_comment(self, column: ColumnInfo) -> Optional[str]:
        """生成基于字段名和类型的备选注释"""
        try:
            column_name = column.column_name.lower()
            data_type = column.data_type.lower()
            
            # 基于字段名的常见模式生成注释
            name_patterns = {
                'id': '唯一标识符',
                'uuid': '全局唯一标识符',
                'name': '名称',
                'title': '标题',
                'desc': '描述',
                'description': '描述信息',
                'content': '内容',
                'text': '文本内容',
                'email': '邮箱地址',
                'phone': '电话号码',
                'mobile': '手机号码',
                'address': '地址',
                'url': 'URL地址',
                'link': '链接',
                'status': '状态',
                'state': '状态',
                'type': '类型',
                'category': '分类',
                'tag': '标签',
                'flag': '标记',
                'enable': '启用状态',
                'disabled': '禁用状态',
                'active': '激活状态',
                'valid': '有效性',
                'code': '编码',
                'no': '编号',
                'num': '数量',
                'count': '计数',
                'total': '总数',
                'amount': '金额',
                'price': '价格',
                'cost': '费用',
                'value': '值',
                'score': '分数',
                'level': '级别',
                'rank': '排名',
                'order': '排序',
                'sort': '排序',
                'priority': '优先级',
                'weight': '权重',
                'size': '大小',
                'length': '长度',
                'width': '宽度',
                'height': '高度',
                'create': '创建',
                'update': '更新',
                'modify': '修改',
                'delete': '删除',
                'time': '时间',
                'date': '日期',
                'year': '年份',
                'month': '月份',
                'day': '日期',
                'hour': '小时',
                'minute': '分钟',
                'second': '秒',
                'user': '用户',
                'admin': '管理员',
                'role': '角色',
                'permission': '权限',
                'group': '组',
                'department': '部门',
                'company': '公司',
                'organization': '组织',
                'version': '版本',
                'revision': '修订版本',
                'comment': '备注',
                'remark': '备注',
                'note': '注释',
                'memo': '备忘',
            }
            
            # 尝试匹配字段名模式
            for pattern, description in name_patterns.items():
                if pattern in column_name:
                    # 根据数据类型调整描述
                    if data_type in ['int', 'integer', 'bigint', 'smallint', 'tinyint']:
                        if 'id' in pattern:
                            return f"{description}(数字)"
                        elif pattern in ['count', 'num', 'total', 'amount', 'price', 'cost', 'score', 'level', 'rank']:
                            return f"{description}(数值)"
                        else:
                            return f"{description}(整数)"
                    elif data_type in ['varchar', 'char', 'text', 'string']:
                        return f"{description}(文本)"
                    elif data_type in ['datetime', 'timestamp', 'date', 'time']:
                        return f"{description}(时间)"
                    elif data_type in ['decimal', 'float', 'double', 'numeric']:
                        return f"{description}(小数)"
                    elif data_type in ['boolean', 'bool', 'bit']:
                        return f"{description}(布尔值)"
                    else:
                        return description
            
            # 基于数据类型的通用描述
            type_descriptions = {
                'int': '整数字段',
                'integer': '整数字段',
                'bigint': '长整数字段',
                'smallint': '短整数字段',
                'tinyint': '微整数字段',
                'varchar': '变长字符串',
                'char': '定长字符串',
                'text': '文本字段',
                'longtext': '长文本字段',
                'datetime': '日期时间',
                'timestamp': '时间戳',
                'date': '日期',
                'time': '时间',
                'decimal': '小数字段',
                'float': '浮点数',
                'double': '双精度浮点数',
                'boolean': '布尔字段',
                'bool': '布尔字段',
                'bit': '位字段',
                'json': 'JSON数据',
                'blob': '二进制数据',
                'binary': '二进制字段',
                'enum': '枚举值',
                'set': '集合值'
            }
            
            # 尝试根据数据类型生成描述
            for type_key, description in type_descriptions.items():
                if type_key in data_type:
                    return description
            
            # 默认描述
            return f"{column.column_name}字段"
            
        except Exception as e:
            self.logger.error(f"生成备选注释失败: {str(e)}")
            return None
    
    def batch_enhance_comments(self, tables: List[TableInfo], 
                             comment_type: str = "both",
                             max_workers: int = 3) -> List[TableInfo]:
        """
        批量增强注释
        
        Args:
            tables: 表列表
            comment_type: 注释类型 ("table", "column", "both")
            max_workers: 最大并发数
        """
        enhanced_tables = []
        
        with ThreadPoolExecutor(max_workers=max_workers) as executor:
            future_to_table = {}
            
            for table in tables:
                future = executor.submit(
                    self._enhance_single_table,
                    table, tables, None, None,
                    comment_type in ["table", "both"],
                    comment_type in ["column", "both"],
                    False
                )
                future_to_table[future] = table
            
            for future in as_completed(future_to_table):
                table = future_to_table[future]
                try:
                    enhanced_table = future.result()
                    enhanced_tables.append(enhanced_table)
                except Exception as e:
                    self.logger.error(f"批量增强失败 {table.table_name}: {str(e)}")
                    enhanced_tables.append(table)
        
        return enhanced_tables
    
    def get_enhancement_statistics(self, schema: DatabaseSchema) -> Dict[str, Any]:
        """获取增强统计信息"""
        stats = {
            "total_tables": 0,
            "enhanced_table_comments": 0,
            "total_columns": 0,
            "enhanced_column_comments": 0,
            "enhancement_rate": 0
        }
        
        for schema_name, tables in schema.schemas.items():
            for table in tables:
                stats["total_tables"] += 1
                if table.table_comment and table.table_comment.strip():
                    stats["enhanced_table_comments"] += 1
                
                for column in table.columns:
                    stats["total_columns"] += 1
                    if column.column_comment and column.column_comment.strip():
                        stats["enhanced_column_comments"] += 1
        
        if stats["total_tables"] > 0:
            table_rate = stats["enhanced_table_comments"] / stats["total_tables"]
            column_rate = stats["enhanced_column_comments"] / stats["total_columns"] if stats["total_columns"] > 0 else 0
            stats["enhancement_rate"] = (table_rate + column_rate) / 2 * 100
        
        return stats


# 全局结构增强器实例
schema_enhancer = SchemaEnhancer()